package com.campus.counseling.config;

import com.campus.counseling.security.JwtTokenUtil;
import com.campus.counseling.security.SecurityUser;
import lombok.RequiredArgsConstructor;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.stereotype.Component;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.security.Principal;

@Component
@RequiredArgsConstructor
public class WebSocketAuthInterceptor implements ChannelInterceptor {

    private final JwtTokenUtil jwtTokenUtil;
    private static final Logger log = LoggerFactory.getLogger(WebSocketAuthInterceptor.class);

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

        if (accessor == null) {
            log.warn("无法获取 StompHeaderAccessor");
            return message;
        }

        if (accessor.getUser() != null) {
            return message;
        }

        if (StompCommand.CONNECT.equals(accessor.getCommand())) {
            String token = accessor.getFirstNativeHeader("Authorization");
            log.debug("收到连接请求，token: {}", token);

            if (token != null) {
                // 如果token不是以Bearer 开头，添加前缀
                if (!token.startsWith("Bearer ")) {
                    token = "Bearer " + token;
                }
                token = token.substring(7);

                try {
                    Authentication auth = jwtTokenUtil.getAuthenticationFromToken(token);
                    if (auth != null && auth.getPrincipal() instanceof SecurityUser) {
                        SecurityUser user = (SecurityUser) auth.getPrincipal();
                        accessor.setUser(new Principal() {
                            @Override
                            public String getName() {
                                return user.getId().toString();
                            }
                        });
                        log.info("WebSocket认证成功 - 用户ID: {}", user.getId());
                        return message;
                    } else {
                        log.warn("无效的认证信息");
                    }
                } catch (Exception e) {
                    log.error("解析token失败: {}", e.getMessage());
                }
            } else {
                log.warn("未找到Authorization header");
            }
            return null; // 认证失败时阻止连接
        }

        // 对于非CONNECT消息，从SecurityContext获取认证信息
        Authentication auth = SecurityContextHolder.getContext().getAuthentication();
        if (auth != null && auth.isAuthenticated()) {
            accessor.setUser(auth);
            accessor.setHeader("user", auth);
            log.debug("从SecurityContext恢复认证信息: {}", auth.getName());
        } else {
            log.warn("无法从SecurityContext获取认证信息");
        }

        return message;
    }
} 